#pragma once

#include "exception.h"
#include <string>
#include <memory>
#include <vector>
#include <utility>

struct LowercaseLatinTraits {
    typedef char Character;
    typedef std::string String;
    static const size_t CharacterCount = 'z' - 'a' + 1;
    static unsigned toLabel(Character c) {
        ASSERT('a' <= c && c <= 'z');
        return c - 'a';
    }
    static Character toCharacter(unsigned label) {
        ASSERT(label < CharacterCount);
        return 'a' + label;
    }
};

struct DefaultTraits {
    typedef char Character;
    typedef std::string String;
    static const size_t AlphabetSize = 'z' - 'a' + 1;
    static const size_t CharacterCount = AlphabetSize +
                                         10 +
                                         1;
    static unsigned toLabel(Character c) {
        if ('a' <= c && c <= 'z')
            return c - 'a';
        else if ('0' <= c && c <= '9')
            return AlphabetSize + c - '0';
        else {
            ASSERT(c == '\'');
            return CharacterCount - 1;
        }
    }
    static Character toCharacter(unsigned label) {
        ASSERT(label < CharacterCount);
        if (label < AlphabetSize)
            return 'a' + label;
        else if (label == CharacterCount - 1)
            return '\'';
        else
            return '0' + (label - AlphabetSize);
    }
};

template <typename WordInfo,
          typename Traits = DefaultTraits>
class Trie {
    typedef typename Traits::String String; 
    class Node;
public:
    Trie()
        : root_(new Node(SpecialLabel))
    {}

    std::pair<WordInfo*, bool> addWord(const String& word) {
        typedef typename String::const_iterator Iterator;
        Node* curNode = root();
        for (Iterator it = word.begin();
             it != word.end();
             ++it)
        {
            unsigned childLabel = Traits::toLabel(*it);
            Node* child = curNode->child(childLabel);
            if (child == 0) {
                curNode = curNode->addChild(childLabel);
            } else {
                curNode = child;
            }
        }
        if (curNode->wordInfo() != 0) {
            return std::make_pair(curNode->wordInfo(), false);
        } else {
            curNode->setWordInfo(WordInfo());
            return std::make_pair(curNode->wordInfo(), true);
        }
    }

    bool contains(const String& word) const {
        const Node* node = find(word);
        return node != 0 && node->wordInfo() != 0;
    }
    
    const WordInfo* operator[](const String& word) const {
         const Node* node = find(word);
         return (node == 0) ? 0 : node->wordInfo();
    }

    WordInfo* operator[](const String& word) {
        Node* node = const_cast<Node*>(find(word));
        return (node == 0) ? 0 : node->wordInfo();
    }

    const Node* root() const {
        return root_.get();
    }

private:
    const Node* find(const typename Traits::String& word) const {
        typedef typename Traits::String::const_iterator Iterator;
        const Node* curNode = root();
        for (Iterator it = word.begin();
             it != word.end() && curNode != 0;
             ++it)
        {
            curNode = curNode->child(Traits::toLabel(*it));
        }
        return curNode;
    }

    Node* root() {
        return root_.get();
    }

private:
    std::auto_ptr<Node> root_;
    static const unsigned SpecialLabel = Traits::CharacterCount;
};

template <typename WordInfo, typename Traits>
class Trie<WordInfo, Traits>::Node {
public:
    explicit Node(unsigned label)
        : label_(label),
          children_(Traits::CharacterCount, 0)
    {
        ASSERT(label_ < children_.size() || label_ == Trie::SpecialLabel);
    }

    ~Node() {
        for (size_t i = 0; i < children_.size(); ++i) {
            delete children_[i];
        }
    }
 
    Node* addChild(unsigned label) {
        ASSERT(label < children_.size()); 
        ASSERT(children_[label] == 0);
        children_[label] = new Node(label);
        return children_[label];
    }

    const Node* child(unsigned label) const {
         return children_.at(label);
    }

    Node* child(unsigned label) { 
         return children_.at(label);
    }
    
    typename Traits::Character character() const {
        ASSERT(!isSpecial());
        return Traits::toCharacter(label_);
    }

    unsigned label() const {
        return label_;
    }

    void setWordInfo(const WordInfo& wordInfo) {
        wordInfo_.reset(new WordInfo(wordInfo));
    }

    const WordInfo* wordInfo() const {
        return wordInfo_.get();
    }
    
    WordInfo* wordInfo() {
        return wordInfo_.get();
    }

    bool isSpecial() const {
        return label_ == Trie::SpecialLabel; 
    }
    
private:
    unsigned label_;
    std::vector<Node*> children_;
    std::auto_ptr<WordInfo> wordInfo_;
};

